#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-09-23


import threading
import time
import cv2
import os
import json
from queue import Queue
from spirems import Publisher, Subscriber, cvimg2sms, sms2cvimg
from spirecv.base.BaseNode import BaseNode
import ffmpeg


class RTMPStreamerNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file = None
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file
        )

        self.image_sms_url = self.get_param("image_sms_url", "/live/sensor/image_raw")
        self.stream_url = self.get_param("stream_url", "rtmp://monitor-push.byaero.com/boyingLive/BY221423030028")
        self.resize = self.get_param("resize", [640, 480])

        self._image_reader = Subscriber(
            self.image_sms_url, 'sensor_msgs::CompressedImage', self.image_callback,
            ip=ip, port=port
        )

        self.image_queue = Queue()
        self.queue_pool.append(self.image_queue)

        self.process = (
            ffmpeg
            .input('pipe:0', format='rawvideo', pix_fmt='bgr24', s='{}x{}'.format(self.resize[0], self.resize[1]), framerate=20)
            .output(self.stream_url, format='flv')
            .run_async(pipe_stdin=True)
        )
        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()

    def image_callback(self, msg):
        self.image_queue.put(msg)

    def run(self):
        while self.is_running():
            msg1 = self.image_queue.get(block=True)
            if msg1 is None:
                break

            img = sms2cvimg(msg1)
            self.process.stdin.write(img.tobytes())
            cv2.imshow('img', img)
            cv2.waitKey(5)

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parameter_file = '../params/default_params.json'
    cam = RTMPStreamerNode("live", param_dict_or_file=parameter_file)
    cam.join()
